package com.it.data_tree.tree.avlTree;

import com.it.data_tree.util.ColorPrintln;
import com.it.data_tree.util.PrintTree;

/**
 * @Author: CaoYouGen
 * @DateTime: 2023/04/27/16:10
 * @注释: AVL树
 **/
public class AVLTree {
    public AVLNode root;
    
    
    /**
     * @param key:
     * @return AVLNode
     * @作者: Coke
     * @注释: 新增和更新节点
     * @DateTime: 2023/5/1 17:22
     */
    public AVLNode putNode (int key) {
        root = doPutNode(root, key);
        return root;
    }
    
    
    private AVLNode doPutNode (AVLNode node, int key) {
        // 1. 找到了空位，创建新节点
        if (node == null) {
            return new AVLNode(key);
        }
        // 2. key已经存在，更新
        if (key == node.val) {
            node.val = key;
            return node;
        }
        // 3. 继续查找
        if (key < node.val) {
            // 向左
            node.left = doPutNode(node.left, key);
        } else {
            // 向右
            node.right = doPutNode(node.right, key);
        }
        // 更新节点高度
        updateHeight(node);
        ColorPrintln.bluePrintln("*************** 插入节点" + key + "旋转前树的结构如下：***************");
        PrintTree.show(root);
        // 平衡树
        return balance(node);
    }
    
    /**
     * @param key:
     * @return AVLNode
     * @作者: Coke
     * @注释: 删除节点
     * @DateTime: 2023/5/1 17:42
     */
    public AVLNode remove (int key) {
        root = doRemove(root, key);
        return root;
    }
    
    private AVLNode doRemove (AVLNode node, int key) {
        // 1. 节点为空
        if (node == null) {
            return null;
        }
        // 2. 没有找到
        if (key < node.val) { // 向左找
            node.left = doRemove(node.left, key);
        } else if (key > node.val) { // 向右找
            node.right = doRemove(node.right, key);
        } else { // 找到了
            // 3.1 叶子节点没有子节点
            if (node.left == null && node.right == null) {
                return null;
                // 3.2 只有一个孩子 (右孩子)
            } else if (node.left == null) {
                node = node.right;
                // 3.3 只有一个孩子 (左孩子)
            } else if (node.right == null) {
                node = node.left;
                // 3.4 有两个孩子
            } else {
                // 找后继节点
                AVLNode rm = node.right;
                while (rm.left != null) {
                    rm = rm.left;
                }
                // 处理后继节点的后事
                rm.right = doRemove(node.right, rm.val);
                rm.left = node.left;
                node = rm;
            }
        }
        // 4. 更新高度
        updateHeight(node);
        // 5. 平衡树
        return balance(node);
    }
    
    /**
     * @param node: 传入节点
     * @return int 返回该节点的高度
     * @Author: CaoYouGen
     * @注释: 获取节点的高度
     * @DateTime: 2023/4/28 9:38
     */
    public int getHeight (AVLNode node) {
        if (node == null) {
            return 0;
        }
        return node.height;
    }
    
    /**
     * @param node:
     * @return void
     * @Author: CaoYouGen
     * @注释: 更新节点的高度
     * @DateTime: 2023/4/28 12:09
     */
    public void updateHeight (AVLNode node) {
        node.height = Integer.max(getHeight(node.left), getHeight(node.right)) + 1;
    }
    
    /**
     * @param node: 传入的节点
     * @return int  平衡因子
     * @Author: CaoYouGen
     * @注释: 获取节点的平衡因子
     * @DateTime: 2023/4/28 9:51
     */
    public int getBalanceFactor (AVLNode node) {
        if (node == null) {
            return 0;
        }
        // 左子树的高度 - 右子树的高度 = 平衡因子
        return getHeight(node.left) - getHeight(node.right);
    }
    
    
    /**
     * @param y: 传入的节点
     * @return AVLNode 新的根节点
     * @Author: CaoYouGen
     * @注释: avl树右旋
     * @DateTime: 2023/4/28 11:30
     */
    // 定义一个函数，用于执行AVL树的右旋操作，输入参数为当前节点y，返回值为新的根节点
        /*
               y5
              / \
             x2  7
            / \
           1   z3

           x2
          / \
        1   y5
          /   \
        z3      7
     */
    /*
首先，我们定义了一个函数rightRotate，输入参数为当前节点y，表示我们要对以y为根节点的子树进行右旋操作。返回值为新的根节点，因为经过右旋操作后，原来的根节点会被替换为右旋后的子树的根节点。

接下来，我们按照上文所述的步骤进行右旋操作。

1. 首先，我们定义一个临时变量x，将当前节点y的左子树保存在x中。

```
Node x = y.left;
```

2. 接着，我们将当前节点y的左子树设置为临时变量x的右子树。这一步的目的是将x的右子树接到y的左子树上。

```
Node temp = x.right;
x.right = y;
```

3. 然后，我们将临时变量x的右子树设置为当前节点y，这一步是将y节点作为x的右子树。

```
y.left = temp;
```

4. 接下来，我们更新当前节点y和临时变量x的高度。因为右旋操作只会影响到节点y和x，所以只需要更新这两个节点的高度即可。

```
y.height = Math.max(getHeight(y.left), getHeight(y.right)) + 1;
x.height = Math.max(getHeight(x.left), getHeight(x.right)) + 1;
```

5. 最后，我们返回新的根节点x。

```
return x;
```

在实现AVL树右旋操作时，需要注意节点高度的更新。AVL树的平衡因子是节点的左子树高度减去右子树高度的差，因此在进行旋转操作后，需要重新计算节点的高度，以保证AVL树的平衡性。

除此之外，还需要注意左旋和右旋操作的顺序，以及每个节点的指针的修改。

    */
    public AVLNode rightRotate (AVLNode y) {
        // 第一步 保存当前节点的左子树
        AVLNode x = y.left;
        // 第二步：将当前节点设置为x的右子树
        AVLNode temp = x.right;
        x.right = y;
        
        // 第三步：将当前节点的左子树设置为temp
        y.left = temp;
        
        // 第四步：更新当前节点和当前节点左子树的高度
        updateHeight(y);
        updateHeight(x);
        
        // 第五步：返回新的根节点
        return x;
    }
    
    /**
     * @param y: 需要左旋的根节点
     * @return AVLNode  旋转后的根节点
     * @Author: CaoYouGen
     * @注释: 左旋操作
     * @DateTime: 2023/4/28 14:47
     */
     /*
               y5
             /   \
            2    x7
                / \
              6   z8

           x7
          /  \
        y5   z8
       / \
      3   6
     */
    public AVLNode leftRotate (AVLNode y) {
        // 第一步：保存当前节点的右子树
        AVLNode x = y.right;
        
        // 第二步：将当前节点设置为x的左子树
        AVLNode temp = x.left;
        x.left = y;
        
        // 第三步：将当前节点的右子树设置为temp
        y.right = temp;
        // 第四步：更新当前节点和当前节点左子树的高度
        updateHeight(y);
        updateHeight(x);
        // 第五步：返回新的根节点
        return x;
    }
    
    /**
     * @param node: 需要左右旋的节点
     * @return AVLNode 返回旋转后的节点
     * @作者: Coke
     * @注释: 左右旋操作
     * @DateTime: 2023/5/1 16:27
     */
    public AVLNode leftRightRotate (AVLNode node) {
        // 第一步：将节点的左节点进行左旋操作
        node.left = leftRotate(node.left);
        // 第二步：将节点进行右旋操作
        return rightRotate(node);
    }
    
    /**
     * @param node: 需要右左旋的节点
     * @return AVLNode 返回右左旋后的节点
     * @作者: Coke
     * @注释: 右左旋操作
     * @DateTime: 2023/5/1 16:31
     */
    public AVLNode rightLeftRotate (AVLNode node) {
        // 第一步: 将节点的右节点进行右旋操作
        node.right = rightRotate(node.right);
        // 第二步：将节点进行左旋操作
        return leftRotate(node);
    }
    
    /**
     * @param node:
     * @return AVLNode
     * @作者: Coke
     * @注释: 检查节点是否失衡，重新平衡代码
     * @DateTime: 2023/5/1 16:46
     */
    public AVLNode balance (AVLNode node) {
        // 如果节点为空，直接返回
        if (node == null) {
            return null;
        }
        // 该节点的平衡因子
        int h = getBalanceFactor(node);
        // 左子节点的平衡因子
        int lh = getBalanceFactor(node.left);
        // 右字节点的平衡因子
        int rh = getBalanceFactor(node.right);
        
        // 左左
        if (h > 1 && lh >= 0) {
            return rightRotate(node);
            // 左右
        } else if (h > 1 && lh < 0) {
            return leftRightRotate(node);
            // 右左
        } else if (h < - 1 && rh > 0) {
            return rightLeftRotate(node);
            // 右右
        } else if ((h < - 1 && rh <= 0)) {
            return leftRotate(node);
        }
        return node;
    }
    
}
